import { defineStore } from 'pinia'
import { ref, computed } from 'vue'

export const useTrainingStore = defineStore('training', () => {
  // State
  const step = ref(0)
  const replicaGroups = ref({})
  const servers = ref({})
  const communicationStatus = ref('stable')
  const rebuildingCommGroup = ref(false)
  const commConnections = ref({})

  // 初始化通信连接
function initCommConnections() {
  const connections = {}
  const allGpus = Object.values(replicaGroups.value)
    .flatMap(group => Object.values(group.gpus))
  
  // 为每个活跃GPU创建连接
  allGpus.filter(gpu => gpu.status === 'active').forEach(gpu => {
    connections[gpu.id] = {
      bandwidths: {},
      status: 'connected'
    }
    
    // 随机生成一些连接带宽(5-40 Gbps)
    allGpus.filter(other => other.id !== gpu.id && other.status === 'active')
      .forEach(other => {
        if (Math.random() > 0.3) { // 70%概率创建连接
          const bw = Math.floor(Math.random() * 35) + 5
          connections[gpu.id].bandwidths[other.id] = bw
        }
      })
  })
  
  commConnections.value = connections
}
  
  // Initialize with sample data
  function initialize() {
    // Create 3 replica groups by default
    for (let i = 1; i <= 3; i++) {
      addReplicaGroup(i)
    }
    
    // Start step counter
    startStepCounter()

    initCommConnections()
  }

  // 更新通信状态当GPU状态变化时
function updateCommConnections() {
  initCommConnections()
}
  
  // Add a new replica group
  function addReplicaGroup(groupId) {
    const gpus = {}
    const serverCount = Math.ceil(Math.random() * 3) // 1-3 servers per group
    
    for (let s = 0; s < serverCount; s++) {
      const serverId = `server-${groupId}-${s}`
      servers.value[serverId] = servers.value[serverId] || { gpus: {} }
      
      const gpuCount = Math.min(8, Math.ceil(Math.random() * 4) + 2) // 3-6 GPUs per server
      
      for (let g = 0; g < gpuCount; g++) {
        const gpuId = `${groupId}-${s}-${g}`
        gpus[gpuId] = {
          id: gpuId,
          server: serverId,
          status: 'active',
          utilization: Math.floor(Math.random() * 30) + 50, // 50-80%
          temperature: Math.floor(Math.random() * 20) + 50, // 50-70°C
          batch: 32,
          lr: 0.001,
          opt: 'Adam'
        }
        servers.value[serverId].gpus[gpuId] = gpus[gpuId]
      }
    }
    
    replicaGroups.value[groupId] = {
      id: groupId,
      gpus,
      status: 'training',
      stepProgress: 0
    }
  }
  
async function deactivateGPU(gpuId) {
  const [groupId] = gpuId.split('-')
  const group = replicaGroups.value[groupId]
  
  // 保存当前step状态
  group.fixedStep = step.value  // 明确记录停止时的step
  group.lastActiveStep = step.value
  group.status = 'deactivating'
  
  // 模拟延迟
  await new Promise(resolve => setTimeout(resolve, 800))
  
  // Mark all GPUs in this replica group as offline
  for (const gpu of Object.values(group.gpus)) {
    gpu.status = 'offline'
    gpu.utilization = 0  // 利用率降为0
    gpu.temperature = 40 // 温度降到基础值
  }
  
  group.status = 'offline'
  communicationStatus.value = 'rebuilding'
  rebuildingCommGroup.value = true
  
  setTimeout(() => {
    communicationStatus.value = 'stable'
    rebuildingCommGroup.value = false
  }, 2000)

  updateCommConnections()
}

// Reactivate a replica group
async function reactivateGroup(groupId) {
  const group = replicaGroups.value[groupId]
  
  // 清除fixedStep，恢复step同步
  group.lastActiveStep = step.value
  delete group.fixedStep
  group.status = 'activating'
  
  // 模拟延迟
  await new Promise(resolve => setTimeout(resolve, 800))
  
  // 恢复所有GPU
  for (const gpu of Object.values(group.gpus)) {
    gpu.status = 'active'
    gpu.utilization = Math.floor(Math.random() * 30) + 50
    gpu.temperature = Math.floor(Math.random() * 20) + 50
  }
  
  group.status = 'training'
  communicationStatus.value = 'rebuilding'
  rebuildingCommGroup.value = true
  
  setTimeout(() => {
    communicationStatus.value = 'stable'
    rebuildingCommGroup.value = false
  }, 2000)

  updateCommConnections()
}

// 修改startStepCounter函数
function startStepCounter() {
  setInterval(() => {
    // 全局step总是增加
    step.value++
    
    // 只更新处于training状态的replica group的step
    Object.values(replicaGroups.value).forEach(group => {
      if (group.status === 'training') {
        // 如果group刚刚被激活，同步到全局step
        if (group.lastActiveStep === undefined) {
          group.stepProgress = step.value % 100
          group.lastActiveStep = step.value
        } 
        // 否则正常增加
        else {
          group.stepProgress = (group.stepProgress + 1) % 100
          group.lastActiveStep = step
        }
      }
    })
  }, 1500)
}
  
  // Computed properties
  const activeGPUs = computed(() => {
    return Object.values(replicaGroups.value)
      .flatMap(group => Object.values(group.gpus))
      .filter(gpu => gpu.status === 'active')
  })
  
  const activeServers = computed(() => {
    const serverIds = new Set()
    activeGPUs.value.forEach(gpu => serverIds.add(gpu.server))
    return Array.from(serverIds)
  })
  
  const activeReplicaGroups = computed(() => {
    return Object.values(replicaGroups.value)
      .filter(group => group.status === 'training')
  })
  
  return {
    step,
    replicaGroups,
    servers,
    communicationStatus,
    rebuildingCommGroup,
    activeGPUs,
    activeServers,
    activeReplicaGroups,
    initialize,
    addReplicaGroup,
    deactivateGPU,
    reactivateGroup,
    startStepCounter
  }
})